import os
from pprint import pprint

from zhipuai import ZhipuAI, APIRequestFailedError

client = ZhipuAI(api_key=os.getenv('ZHIPUAI_API_KEY'))  # 请填写您自己的APIKey

def zhipu_call(model="GLM-4-Air-250414", convs=None):
    if convs is None:
        convs = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "中国的首都是是哪？网上查一下，然后简短的告诉我。"}
        ]
    response = client.chat.completions.create(
        model=model,
        messages=convs,
    )
    return_body = {
        "search": [],
        "thinking": response.choices[0].message.reasoning_content,
        "answer": response.choices[0].message.content,
        "usage": {
            "input_tokens": response.usage.prompt_tokens,
            "output_tokens": response.usage.completion_tokens,
            "total_tokens": response.usage.total_tokens
        }
    }
    return return_body

if __name__ == '__main__':
    res = zhipu_call()
    pprint(res)
